//============================================================================
// Name        : csiena
// Author      : Josh Lospinoso
// Copyright   : MIT License
// Description : Austere, limited implementation of the R package RSiena in
//				 C++. Used for fast maximum likelihood and method of moments
//				 estimation of stochastic actor oriented models (SAOMs)
//				 without any frills or user experience niceties.
//============================================================================

#include "digraph.h"
#include "link.h"
#include "globals.h"
#include "evaluation.h"
#include "chain.h"
#include "mersenneTwister.h"

#ifndef MCMCMLE_H_
#define MCMCMLE_H_
namespace csiena {

class mcmcmle {
public:
	// Initializers
	mcmcmle(int actors, int seed);
	virtual ~mcmcmle();

	// Accessor methods
	/* Theta */
	void setBeta(int i, double beta);
	void setGamma(int i, double gamma);
	void setAlpha(double alpha);
	double getAlpha();
	double getBeta(int i);
	double getGamma(int i);
	/* Likelihood */
	double getLogLikelihood();
	/* Scores */
	double getScore(int k);
	double getGammaScore(int k);
	double getBetaScore(int k);
	double getAlphaScore();
	/* Information */
	double getAlphaInformation();
	double getBetaInformation(int j, int k);
	double getGammaInformation(int j, int k);
	double getInformation(int k, int l);
	double getObservedInformation(int j, int k);
	double getScoreCovariance(int j, int k);
	double getStandardErrorSquared(int i, int j);
	/* Other */
	int getActors();
	int getSampleSize();
	double getAverageSubphaseTheta(int i);
	double getSubphaseScoreAutocorr(int i);
	double getProposalAcceptanceRatio();
	double getConvergenceStatistic(int i);

	// Utility methods
	void print();
	void print(bool verbose);
	void loadStart(const char* fname);
	void loadObservation(const char* fname);

	// Robbins-Monroe Related
	void estimateMCMC(int sampleSpacing, int samples, bool estimateInformation);
	void estimateCompleteData();
	void partialNewtonRaphsonStep(double gain);
	void partialSingleCoordinateNewtonRaphsonStep(double gain);
	void preprocessDerivative();
	void executePhaseOne(bool singleCoordinate, bool useCompleteInformation,
			bool diagonalOnly, int burns, int sampleSpacing, int samples,
			double gain, bool verbose);
	void executePhaseTwoSubphase(int iterationsMin, int iterationsMax, double gain,
			bool singleCoordinate, int sampleSpacing, int samples, bool verbose);
	void executePhaseThree(int sampleSpacing, int samples, bool verbose);
	// MCMC Related
	void generateSimulatedChain(int chainSize, int burnInsForStart, int observations);
	void generateSimulatedChain(int chainSize);
	void burnChain(int draws);

	// Diagnostic methods
	void estimateProfileLogLikelihood(int betaIndex, double betaMin, double betaMax,
			int samples, double* betaFixed, double* profileLogLikelihood,
			int MCMCsamples, int sampleSpacing);

	// Ingredients
	Chain *completeData;
	MersenneTwister *r;

	double getDerivative(int i, int j);
	void setDerivative(int i, int j, double value);
private:
	#ifdef BOOST_THREADS /* These are only needed when multithreading */
		class EstimationResults {
		public:
			double logLikelihoodSum;
			double *scoresSum;
			double *scoresSquaredSum;
			double *informationSum;
			int mcmcSamplesDrawn;
			int mcmcProposalsAccepted;
			int mcmcProposalsMade;
			EstimationResults() {
				this->scoresSum = new double[BETA_LENGTH];
				this->scoresSquaredSum = new double[BETA_LENGTH*BETA_LENGTH];
				this->informationSum = new double[BETA_LENGTH*BETA_LENGTH];
			}
			~EstimationResults() {
				delete(this->scoresSum);
				delete(this->scoresSquaredSum);
				delete(this->informationSum);
			}
		};
		class EstimationWorker {
		public:
			EstimationWorker(Chain* assignedChain, MersenneTwister* randomNumGen) :
				threadChain(assignedChain), r(randomNumGen) {

			}
			void run (EstimationResults *results, int sampleSpacing, int samples) {
				// Reset the statistics:
				results->mcmcProposalsMade = 0;
				results->mcmcProposalsAccepted = 0;
				results->logLikelihoodSum = 0;
				results->logLikelihoodSum = 0;
				for (int k=0; k < BETA_LENGTH; k++) {
					results->scoresSum[k] = 0;
					for (int l=0; l<BETA_LENGTH; l++) {
						results->informationSum[k*BETA_LENGTH+l] = 0;
						results->scoresSquaredSum[k*BETA_LENGTH+l] = 0;
					}
				}
				for (int draw=0; draw < samples * (sampleSpacing+1); draw++) {
					results->mcmcProposalsMade++;
					if (threadChain->drawProposal(r)) {
						results->mcmcProposalsAccepted++;
					}
					if (draw % (sampleSpacing+1) == 0) {
						threadChain->evaluateChain();
						// Draw sample
						results->logLikelihoodSum += threadChain->getLogLikelihood();
						for (int k=0; k < BETA_LENGTH; k++) {
							results->scoresSum[k] += threadChain->getBetaScore(k);
							for (int l=0; l<BETA_LENGTH; l++) {
								results->informationSum[k*BETA_LENGTH+l] += threadChain->getBetaInformation(k, l);
								results->scoresSquaredSum[k*BETA_LENGTH+l] += threadChain->getBetaScore(k)*threadChain->getBetaScore(l);
							}
						}
					}
				}
			}
		private:
			Chain* threadChain;
			MersenneTwister* r;
		};
		EstimationWorker *workers[BOOST_THREADS];
		boost::thread *workerThreads[BOOST_THREADS];
		EstimationResults *results[BOOST_THREADS];
		Chain *threadChains[BOOST_THREADS-1];
		MersenneTwister *threadR[BOOST_THREADS];
	#endif /*BOOST_THREADS*/
	void setStandardErrorSquared(int i, int j, double value);

	void calculateDerivativeMatrix();
	void calculateDerivativeMatrix(bool useCompleteInformation, bool diagonalOnly);
	void calculateConvergenceStatistics();
	void calculateStandardErrorsSquared();
	/* Likelihood */
	void setLogLikelihoodSum(double value);
	void incrementLogLikelihoodSum(double value);
	void resetLikelihoodStatistics(bool scoresOnly);

	/* Scores */
	void setScoreSquaredSum(int j, int k, double value);
	double getScoreSquaredSum(int j, int k);
	void setScoreSum(int k, double value);
	void setAlphaScoreSum(double value);
	void setBetaScoreSum(int k, double value);
	void setGammaScoreSum(int k, double value);
	void incrementAlphaScoreSum(double value);
	void incrementBetaScoreSum(int k, double value);
	void incrementGammaScoreSum(int k, double value);
	void incrementScoreSum(int k, double value);
	void incrementScoreSquaredSum(int j, int k, double value);

	/* Information */
	void setAlphaInformationSum(double value);
	void setBetaInformationSum(int j, int k, double value);
	void setInformationSum(int a, int b, double value);
	void setGammaInformationSum(int j, int k, double value);
	void incrementAlphaInformationSum(double value);
	void incrementBetaInformationSum(int j, int k, double value);
	void incrementGammaInformationSum(int j, int k, double value);
	void incrementInformationSum(int j, int k, double value);

	void setConvergenceStatistic(int i, double value);
	int actors;
	int subphaseN;
	int sampleSize;
	double logLikelihood;
	double* betaScore;
	double* betaInformation;
	double* gammaScore;
	double* gammaInformation;
	double alphaScore;
	double alphaInformation;
	double* scoreSquaredSum;
	double* derivativeMatrix;
	double* standardErrorsSquared;
	double* convergenceStatistics;
	double *subphaseSumTheta;
	double *subphasePreviousScore;
	double *subphaseAutocorrScore;
	int chainLengthSum;
	double subphaseSumAlpha;
	double mcmcProposalsAcceptedInPhase;
	int mcmcProposalsMadeInPhase;
	#ifndef CUDA
	gsl_matrix *matrixGSL;
	gsl_matrix *matrixInverseGSL;
	gsl_permutation *luPermutationGSL;
	#else
	culaStatus status;
	#endif
};
}
#endif /* MCMCMLE_H_ */
